import os
import wave

# 遍历指定文件夹下的所有.wav文件
def process_folder(input_folder, output_folder):
    for file_name in os.listdir(input_folder):
        if file_name.endswith(".wav"):
            file_path = os.path.join(input_folder, file_name)
            process_wav_file(file_path, output_folder)
            print(file_name)

# 处理单个.wav文件
def process_wav_file(file_path, output_folder):
    # 读取.wav文件
    with wave.open(file_path, 'rb') as wav_file:
        sample_rate = wav_file.getframerate()
        num_samples = wav_file.getnframes()
        num_channels = wav_file.getnchannels()
        sample_width = wav_file.getsampwidth()

        # 将每个.wav文件分为等长的n段
        n = 6
        segment_size = int(num_samples / n)
        for i in range(n):
            start_sample = i * segment_size
            end_sample = start_sample + segment_size
            if i == n-1:
                end_sample = num_samples  # 最后一段包含剩下所有的样本
            segment_data = wav_file.readframes(end_sample - start_sample)

            # 将分段后的音频保存到对应的文件夹中
            output_subfolder = os.path.splitext(file_path)[0]
            output_folder_path = os.path.join(output_folder, output_subfolder)
            if not os.path.exists(output_folder_path):
                os.makedirs(output_folder_path)
            output_file_name = f"{i+1}.wav"
            output_file_path = os.path.join(output_folder_path, output_file_name)
            with wave.open(output_file_path, 'wb') as output_file:
                output_file.setnchannels(num_channels)
                output_file.setsampwidth(sample_width)
                output_file.setframerate(sample_rate)
                output_file.writeframes(segment_data)

# 在这里指定要遍历的文件夹和输出文件夹
input_folder = "music_dataset"     # music_dataset
output_folder = "Emo/video"   # music_segment_dataset_5
process_folder(input_folder, output_folder)
